package com.ruoyi.util;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.text.PDFTextStripper;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.List;
import java.util.Set;
import java.util.*;
public class PdfPartNumberProcessor {
    private final Set<String> partNumbers;
    private final String outputDirectory;

    public PdfPartNumberProcessor(List<String> partNumbers, String outputDirectory) {
        // 转换为小写的HashSet以提高查询效率
        this.partNumbers = partNumbers.stream()
                .map(String::toLowerCase)
                .collect(java.util.stream.Collectors.toCollection(HashSet::new));
        this.outputDirectory = outputDirectory;
    }

    public Map<String, String> processPdf(InputStream pdfInputStream) throws IOException {
        // 主文档只打开一次，直到所有处理完成后再关闭
        try (PDDocument mainDocument = PDDocument.load(pdfInputStream)) {
            int totalPages = mainDocument.getNumberOfPages();
            Map<String, PDDocument> partNumberDocuments = new HashMap<>();

            try {
                // 初始化每个零件号对应的文档
                for (String partNumber : partNumbers) {
                    partNumberDocuments.put(partNumber, new PDDocument());
                }

                // 逐页处理PDF
                for (int pageNum = 0; pageNum < totalPages; pageNum++) {
                    // 关键修复：创建页面的深层复制，避免依赖主文档的资源
                    PDDocument singlePageDoc = new PDDocument();
                    singlePageDoc.importPage(mainDocument.getPage(pageNum));

                    try {
                        // 提取页面文本
                        String pageText = extractText(singlePageDoc).toLowerCase();

                        // 检查该页面是否包含任何零件号
                        for (String partNumber : partNumbers) {
                            if (pageText.contains(partNumber)) {
                                // 将页面添加到对应的文档中
                                partNumberDocuments.get(partNumber).importPage(singlePageDoc.getPage(0));
                            }
                        }
                    } finally {
                        // 及时临时文档，释放资源
                        singlePageDoc.close();
                    }
                }

                // 保存所有零件号对应的文档
                return saveDocumentResults(partNumberDocuments);
            } finally {
                // 确保所有临时文档都被关闭
                for (PDDocument doc : partNumberDocuments.values()) {
                    if (doc != null) {
                        doc.close();
                    }
                }
            }
        }
    }

    // 提取PDF文本内容
    private String extractText(PDDocument document) throws IOException {
        PDFTextStripper stripper = new PDFTextStripper();
        stripper.setStartPage(1);
        stripper.setEndPage(document.getNumberOfPages());
        return stripper.getText(document);
    }

    // 保存处理结果
    private Map<String, String> saveDocumentResults(Map<String, PDDocument> partNumberDocuments) throws IOException {
        Map<String, String> resultPaths = new HashMap<>();

        for (Map.Entry<String, PDDocument> entry : partNumberDocuments.entrySet()) {
            String partNumber = entry.getKey();
            PDDocument document = entry.getValue();

            // 创建零件号对应的文件夹
            Path partDir = Path.of(outputDirectory, partNumber);
            Files.createDirectories(partDir);

            // 生成输出文件路径
            String fileName = partNumber + ".pdf";
            Path outputPath = partDir.resolve(fileName);

            // 只有当文档包含页面时才保存
            if (document.getNumberOfPages() > 0) {
                document.save(outputPath.toFile());
                resultPaths.put(partNumber, outputPath.toString());
            } else {
                resultPaths.put(partNumber, "未找到匹配的页面");
            }
        }

        return resultPaths;
    }
}
